import os
import urllib.request
import zipfile
import sys

def download_file(url, filename):
    print(f"正在下载: {filename}")
    try:
        urllib.request.urlretrieve(url, filename)
        print(f"下载成功: {filename}")
        return True
    except Exception as e:
        print(f"下载失败: {e}")
        return False

def unzip_file(zip_filename, extract_dir):
    print(f"正在解压: {zip_filename}")
    try:
        with zipfile.ZipFile(zip_filename, 'r') as zip_ref:
            zip_ref.extractall(extract_dir)
        print(f"解压成功: {zip_filename}")
        return True
    except Exception as e:
        print(f"解压失败: {e}")
        return False

def main():
    print("===== 性别和年龄模型下载脚本 ======")

    # 模型文件列表
    models = [
        {
            "name": "gender_age_models",
            "url": "https://github.com/GilLevi/AgeGenderDeepLearning/archive/refs/heads/master.zip",
            "zip_file": "age_gender_models.zip",
            "extract_dir": "."
        }
    ]

    # 下载并解压模型
    for model in models:
        if not download_file(model["url"], model["zip_file"]):
            print(f"无法下载 {model['name']} 模型")
            continue

        if not unzip_file(model["zip_file"], model["extract_dir"]):
            print(f"无法解压 {model['name']} 模型")
            continue

        # 移动模型文件到当前目录
        try:
            import shutil
            src_dir = os.path.join(model["extract_dir"], "AgeGenderDeepLearning-master")
            if os.path.exists(src_dir):
                # 复制 prototxt 和 caffemodel 文件
                for file_name in os.listdir(src_dir):
                    if file_name.endswith(".prototxt") or file_name.endswith(".caffemodel"):
                        src_path = os.path.join(src_dir, file_name)
                        dst_path = os.path.join(model["extract_dir"], file_name)
                        shutil.copy2(src_path, dst_path)
                        print(f"已复制: {file_name}")

                # 清理临时文件
                shutil.rmtree(src_dir)
                os.remove(model["zip_file"])
                print("已清理临时文件")
        except Exception as e:
            print(f"移动模型文件时出错: {e}")

    print("\n===== 模型下载完成 ======")
    print("请检查当前目录是否包含以下文件:")
    print("- deploy_gender.prototxt")
    print("- gender_net.caffemodel")
    print("- deploy_age.prototxt")
    print("- age_net.caffemodel")
    print("\n如果文件存在，您可以重新运行人脸识别应用程序。")

if __name__ == "__main__":
    main()